import sys
import asyncio
import json
import threading
import traceback
from PyQt5.QtWidgets import (QApplication, QWidget, QLineEdit, QPushButton,
                             QVBoxLayout, QTextEdit, QLabel, QFrame, QHBoxLayout,
                             QListWidget, QListWidgetItem, QAbstractItemView, QStyledItemDelegate,
                             QSizePolicy, QScrollArea)
from PyQt5.QtCore import Qt, QPoint, QTimer, pyqtSignal, QThread, QSize
from PyQt5.QtGui import QFont, QPalette, QColor, QIcon, QTextCharFormat, QTextCursor, QBrush
from mcp_server import chat


class CustomListWidget(QListWidget):
    """Custom list widget to handle item clicks for copy buttons"""
    copyClicked = pyqtSignal(QListWidgetItem)

    def mousePressEvent(self, event):
        super().mousePressEvent(event)
        item = self.itemAt(event.pos())
        if item:
            self.copyClicked.emit(item)


class CommandProcessor(QThread):
    """独立的命令处理线程"""
    resultReady = pyqtSignal(str, str)  # (sender, message)
    statusReady = pyqtSignal(str, str)  # (color, text)

    def __init__(self, command):
        super().__init__()
        self.command = command

    def run(self):
        try:
            self.statusReady.emit("#f39c12", "处理中: 正在连接AI服务...")
            asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

            # 通知正在等待AI响应
            self.statusReady.emit("#3498db", "处理中: 等待AI响应...")

            # 创建事件循环
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            result = loop.run_until_complete(chat(self.command))

            # 处理换行符
            result = result.replace('\n', '<br>')

            # 检查是否是从工具返回的结果
            if "调用工具返回" in result:
                self.statusReady.emit("#3498db", "处理中: MCP工具输出中...")

            self.resultReady.emit("AI", result)
            loop.close()
        except Exception as e:
            error_trace = traceback.format_exc()
            result = f"运行错误: {str(e)}<br>{error_trace}".replace('\n', '<br>')
            self.resultReady.emit("系统", result)
        finally:
            self.statusReady.emit("#2ecc71", "就绪")


class FloatingWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.setupUI()

        # 状态初始化
        self.status = "就绪"
        self.setFixedSize(300, 70)  # 更小的初始尺寸
        self.message_history = []

        # 窗口位置控制
        self.setGeometry(
            QApplication.primaryScreen().availableGeometry().width() - 330,
            50,
            300,
            70
        )

        # 设置窗口透明度
        self.setWindowOpacity(0.95)

    def setupUI(self):
        # 设置窗口属性 - 最小化UI框架
        self.setWindowTitle("AI助手")
        self.setWindowFlags(
            Qt.FramelessWindowHint |
            Qt.WindowStaysOnTopHint |
            Qt.Tool
        )
        self.setStyleSheet("""
            QWidget {
                background-color: rgba(44, 62, 80, 0.8);
                border-radius: 8px;
                padding: 5px;
            }
            QLineEdit {
                background-color: rgba(52, 73, 94, 0.9);
                color: #ecf0f1;
                border: 1px solid #3498db;
                border-radius: 4px;
                padding: 5px;
                font-size: 14px;
            }
            QPushButton {
                background-color: rgba(52, 152, 219, 0.9);
                color: white;
                border: none;
                border-radius: 4px;
                padding: 5px;
                min-width: 50px;
                font-size: 13px;
            }
            QPushButton:hover {
                background-color: rgba(41, 128, 185, 0.9);
            }
            QPushButton:pressed {
                background-color: rgba(21, 67, 96, 0.9);
            }
            QStatusLabel {
                color: #ecf0f1;
                font-size: 12px;
                padding: 2px;
                background-color: rgba(30, 40, 50, 0.7);
                border-radius: 4px;
            }
        """)

        # 创建简洁布局
        main_layout = QVBoxLayout()
        main_layout.setContentsMargins(8, 8, 8, 8)
        main_layout.setSpacing(8)

        # 状态栏 (显示在上方)
        status_layout = QHBoxLayout()
        self.status_icon = QLabel("🟢")
        self.status_label = QLabel("就绪")
        self.status_label.setObjectName("StatusLabel")
        self.status_label.setStyleSheet("color: #2ecc71;")
        self.status_label.setAlignment(Qt.AlignLeft | Qt.AlignVCenter)
        status_layout.addWidget(self.status_icon)
        status_layout.addWidget(self.status_label)
        status_layout.addStretch()
        main_layout.addLayout(status_layout)

        # 输入行布局
        input_layout = QHBoxLayout()
        input_layout.setContentsMargins(0, 0, 0, 0)
        input_layout.setSpacing(8)

        # 输入框
        self.command_entry = QLineEdit()
        self.command_entry.setPlaceholderText("输入命令...")
        self.command_entry.returnPressed.connect(self.send_command)
        input_layout.addWidget(self.command_entry)

        # 发送按钮
        self.send_btn = QPushButton("发送")
        self.send_btn.setFixedWidth(50)
        self.send_btn.clicked.connect(self.send_command)
        input_layout.addWidget(self.send_btn)

        # 关闭按钮
        self.close_btn = QPushButton("×")
        self.close_btn.setFixedSize(25, 25)
        self.close_btn.clicked.connect(self.close)
        input_layout.addWidget(self.close_btn)

        main_layout.addLayout(input_layout)

        self.setLayout(main_layout)

        # 响应窗口
        self.response_window = None

    def send_command(self):
        command = self.command_entry.text().strip()
        if not command:
            return

        # 清空输入框
        self.command_entry.clear()

        # 记录用户消息
        self.add_to_message_buffer("用户", command)
        self.show_response_window()

        # 更新状态
        self.update_status("#f39c12", "处理中: 正在连接...")

        # 在后台线程处理命令
        self.processor = CommandProcessor(command)
        self.processor.resultReady.connect(self.handle_result)
        self.processor.statusReady.connect(self.update_status)
        self.processor.start()

    def handle_result(self, sender, message):
        self.add_to_message_buffer(sender, message)

    def update_status(self, color, text):
        # 更新状态图标
        if "就绪" in text:
            self.status_icon.setText("🟢")
            self.status_label.setStyleSheet(f"color: #2ecc71;")
        elif "处理中" in text:
            self.status_icon.setText("🟡")
            self.status_label.setStyleSheet(f"color: {color};")
        elif "错误" in text:
            self.status_icon.setText("🔴")
            self.status_label.setStyleSheet(f"color: #e74c3c;")

        self.status_label.setText(text)

        # 更新响应窗口状态
        if self.response_window:
            self.response_window.update_status(color, text)

    def add_to_message_buffer(self, sender, message):
        self.message_history.append({"sender": sender, "message": message})

        # 更新响应窗口内容
        if self.response_window:
            self.response_window.add_message(sender, message)

    def show_response_window(self):
        """显示响应窗口（如果未显示）"""
        if not self.response_window or not self.response_window.isVisible():
            self.response_window = ResponseWindow(self)
            # 加载历史消息
            for msg in self.message_history:
                self.response_window.add_message(msg["sender"], msg["message"])
            # 将窗口定位到主窗口下方
            self.response_window.move(
                self.x(),
                self.y() + self.height() + 5
            )
            self.response_window.show()
        else:
            # 确保最新消息可见
            self.response_window.scroll_to_bottom()

    # 窗口拖动功能
    def mousePressEvent(self, event):
        self.drag_start_position = event.globalPos()
        self.window_position = self.pos()

    def mouseMoveEvent(self, event):
        if hasattr(self, 'drag_start_position'):
            delta = event.globalPos() - self.drag_start_position
            self.move(self.window_position + delta)
            # 移动响应窗口（如果有）
            if self.response_window and self.response_window.isVisible():
                self.response_window.move(
                    self.x(),
                    self.y() + self.height() + 5
                )

    def mouseReleaseEvent(self, event):
        if hasattr(self, 'drag_start_position'):
            del self.drag_start_position
            # 保存位置
            self.save_window_position()
            # 更新响应窗口位置（如果有）
            if self.response_window and self.response_window.isVisible():
                self.response_window.move(
                    self.x(),
                    self.y() + self.height() + 5
                )

    def save_window_position(self):
        """保存窗口位置（简化版）"""
        pass


class ResponseWindow(QWidget):
    """单独的响应显示窗口"""

    def __init__(self, parent=None):
        super().__init__(parent)
        self.setupUI()
        self.message_count = 0
        self.status = "就绪"

    def setupUI(self):
        # 设置窗口属性
        self.setWindowTitle("AI响应")
        self.setWindowFlags(
            Qt.FramelessWindowHint |
            Qt.WindowStaysOnTopHint |
            Qt.Tool
        )
        self.setMinimumSize(400, 300)
        self.setStyleSheet("""
            QWidget {
                background-color: rgba(44, 62, 80, 0.95);
                border-radius: 8px;
            }
            QListWidget {
                background-color: rgba(30, 40, 50, 0.8);
                border: none;
                border-radius: 8px;
            }
            QStatusBar {
                color: #ecf0f1;
                font-size: 12px;
                background-color: rgba(30, 40, 50, 0.7);
                border-radius: 0 0 8px 8px;
                padding: 4px;
            }
            QTitleBar {
                background-color: rgba(30, 40, 50, 0.9);
                border-top-left-radius: 8px;
                border-top-right-radius: 8px;
                padding: 4px;
            }
            QCloseBtn {
                background-color: rgba(231, 76, 60, 0.8);
                color: white;
                border-radius: 10px;
                min-width: 25px;
                min-height: 25px;
                font-size: 16px;
            }
            QCloseBtn:hover {
                background-color: rgba(192, 57, 43, 0.8);
            }
        """)

        # 主布局
        main_layout = QVBoxLayout()
        main_layout.setContentsMargins(0, 0, 0, 0)
        main_layout.setSpacing(0)

        # 标题栏 (顶部)
        title_bar = QWidget()
        title_bar.setObjectName("TitleBar")
        title_layout = QHBoxLayout(title_bar)
        title_layout.setContentsMargins(8, 4, 8, 4)
        title_bar.setFixedHeight(35)

        # 标题
        title = QLabel("聊天")
        title.setStyleSheet("color: #ecf0f1; font-weight: bold; font-size: 14px;")
        title_layout.addWidget(title)

        # 弹簧
        title_layout.addStretch()

        # 关闭按钮 (右上角)
        self.close_btn = QPushButton("×")
        self.close_btn.setObjectName("CloseBtn")
        self.close_btn.setFixedSize(25, 25)
        self.close_btn.clicked.connect(self.hide)
        title_layout.addWidget(self.close_btn)

        main_layout.addWidget(title_bar)

        # 创建自定义列表控件
        self.message_list = CustomListWidget()
        self.message_list.setWordWrap(True)
        self.message_list.setVerticalScrollMode(QAbstractItemView.ScrollPerPixel)
        self.message_list.setStyleSheet("""
            QListWidget {
                border-top: 1px solid rgba(255, 255, 255, 0.1);
                border-bottom: 1px solid rgba(255, 255, 255, 0.1);
                background-color: rgba(30, 40, 50, 0.8);
            }
            QListWidget::item {
                padding: 8px;
                border-radius: 4px;
                margin-bottom: 5px;
                border: 1px solid rgba(255, 255, 255, 0.1);
            }
            QListWidget::item:selected {
                background-color: rgba(52, 152, 219, 0.3);
            }
        """)
        self.message_list.copyClicked.connect(self.copy_message)
        main_layout.addWidget(self.message_list)

        # 状态栏 (底部)
        self.status_bar = QLabel("就绪")
        self.status_bar.setObjectName("StatusBar")
        self.status_bar.setAlignment(Qt.AlignCenter)
        self.status_bar.setFixedHeight(25)
        self.status_bar.setStyleSheet("color: #2ecc71;")
        main_layout.addWidget(self.status_bar)

        self.setLayout(main_layout)

    def add_message(self, sender, message):
        # 创建条目和自定义小部件
        item = QListWidgetItem()

        # 根据发送者设置背景色
        if sender == "用户":
            bg_color = "rgba(52, 152, 219, 0.25)"  # 蓝色
            text_color = "#3498db"
            sender_icon = "👤"
        elif sender == "AI":
            bg_color = "rgba(46, 204, 113, 0.25)"  # 绿色
            text_color = "#2ecc71"
            sender_icon = "🤖"
        else:  # 系统消息
            bg_color = "rgba(149, 165, 166, 0.25)"  # 灰色
            text_color = "#95a5a6"
            sender_icon = "⚙️"

        # 添加复制按钮到AI的消息
        copy_btn = ""
        if sender == "AI" or sender == "系统":
            copy_btn = f"""<a href="copy" style="color:#e67e22; text-decoration:none; padding: 2px 5px; background: rgba(230, 126, 34, 0.2); border-radius: 3px;">📋 复制</a>"""

        # 创建条目内容
        content = f"""
        <div style="background-color: {bg_color}; padding: 10px; border-radius: 8px;">
            <div style="font-size: 14px; color: {text_color}; font-weight: bold;">
                {sender_icon} {sender}:
            </div>
            <div style="font-size: 14px; color: #ecf0f1; margin: 10px 0;">
                {message}
            </div>
            <div style="text-align: right; font-size: 12px;">
                {copy_btn}
            </div>
        </div>
        """

        # 创建并设置小部件
        label = QLabel(content)
        label.setOpenExternalLinks(False)
        label.linkActivated.connect(lambda: self.copy_message(item))
        label.setTextInteractionFlags(Qt.TextSelectableByMouse | Qt.LinksAccessibleByMouse)
        label.setWordWrap(True)
        label.setMargin(8)

        # 计算合适的高度
        label.setText(content)  # 确保设置文本
        width = self.message_list.viewport().width() - 30
        height = label.sizeHint().height()
        item.setSizeHint(QSize(width, height + 20))

        # 添加到列表
        self.message_list.addItem(item)
        self.message_list.setItemWidget(item, label)

        # 滚动到底部
        self.scroll_to_bottom()
        self.message_count += 1

    def copy_message(self, item):
        """复制消息到剪贴板"""
        widget = self.message_list.itemWidget(item)
        # 提取消息内容
        if widget:
            # 获取QLabel内部的HTML文本
            html = widget.text()
            # 提取消息文本
            start = html.find('">', html.find('ecf0f1')) + 2
            end = html.find("</div>", start)
            if start > 0 and end > 0:
                content = html[start:end].strip()
                # 将HTML换行符转换回普通换行符
                content = content.replace('<br>', '\n')

                clipboard = QApplication.clipboard()
                clipboard.setText(content)
                # 临时显示复制成功的通知
                self.status_bar.setText("已复制到剪贴板！")
                self.status_bar.setStyleSheet("color: #e67e22;")
                QTimer.singleShot(2000, lambda: self.update_status("#2ecc71", self.status))

    def update_status(self, color, text):
        """更新状态栏文本"""
        self.status = text
        if "就绪" in text:
            self.status_bar.setStyleSheet("color: #2ecc71;")
        elif "处理中" in text:
            self.status_bar.setStyleSheet(f"color: {color};")
        elif "错误" in text:
            self.status_bar.setStyleSheet("color: #e74c3c;")

        self.status_bar.setText(text)

    def scroll_to_bottom(self):
        """滚动到列表底部"""
        if self.message_list.count() > 0:
            self.message_list.scrollToBottom()

    def closeEvent(self, event):
        """关闭时仅隐藏，不销毁"""
        event.ignore()
        self.hide()


if __name__ == "__main__":
    app = QApplication(sys.argv)
    app.setStyle('Fusion')  # 使用更现代的样式

    # 设置应用属性
    app.setApplicationName("PCAgent悬浮窗")

    window = FloatingWindow()
    window.show()

    sys.exit(app.exec_())